{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_func(example):\n",
    "  \"\"\"\n",
    "  将数据集进行预处理, 处理成模型可以接受的格式\n",
    "  \"\"\"\n",
    "\n",
    "  MAX_LENGTH = 384 \n",
    "  input_ids, attention_mask, labels = [], [], []\n",
    "  system_prompt = \"\"\"你是一个文本实体识别领域的专家，你需要从给定的句子中提取 地点; 人名; 地理实体; 组织 实体. 以 json 格式输出, 如 {\"entity_text\": \"南京\", \"entity_label\": \"地理实体\"} 注意: 1. 输出的每一行都必须是正确的 json 字符串. 2. 找不到任何实体时, 输出\"没有找到任何实体\".\"\"\"\n",
    "  \n",
    "  instruction = tokenizer(\n",
    "    f\"<|im_start|>system\\n{system_prompt}<|im_end|>\\n<|im_start|>user\\n{example['input']}<|im_end|>\\n<|im_start|>assistant\\n\",\n",
    "    add_special_tokens=False,\n",
    "  )\n",
    "  response = tokenizer(f\"{example['output']}\", add_special_tokens=False)\n",
    "  input_ids = instruction[\"input_ids\"] + response[\"input_ids\"] + [tokenizer.pad_token_id]\n",
    "  attention_mask = (\n",
    "      instruction[\"attention_mask\"] + response[\"attention_mask\"] + [1]\n",
    "  )\n",
    "  labels = [-100] * len(instruction[\"input_ids\"]) + response[\"input_ids\"] + [tokenizer.pad_token_id]\n",
    "  if len(input_ids) > MAX_LENGTH:  # 做一个截断\n",
    "      input_ids = input_ids[:MAX_LENGTH]\n",
    "      attention_mask = attention_mask[:MAX_LENGTH]\n",
    "      labels = labels[:MAX_LENGTH]\n",
    "  return {\"input_ids\": input_ids, \"attention_mask\": attention_mask, \"labels\": labels}   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import Dataset\n",
    "\n",
    "train_ds = Dataset.from_pandas(train_df)\n",
    "train_dataset = train_ds.map(process_func, remove_columns=train_ds.column_names)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.13.5 (tags/v3.13.5:6cb20a2, Jun 11 2025, 16:15:46) [MSC v.1943 64 bit (AMD64)]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "340a546b2c4c9cf5d23eb4a2a4e78e923e0b7afe4d00162258c4442b3ee3b061"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
